import cv2
import os,shutil
import copy
from glob import glob
import xml.etree.ElementTree as ET
import re
import cv2
import numpy as np
import random

img_format = ['bmp', 'dng', 'jpeg', 'jpg', 'mpo', 'png', 'tif', 'tiff', 'webp']


def drow_rect(img,boxs):
    for i in range(len(boxs)):
        print(boxs[i])
        cv2.rectangle(img,(boxs[i][0],boxs[i][1]),(boxs[i][2],boxs[i][3]),(255,255,0),2)
        cv2.putText(img, boxs[i][4] + ':' + f"{boxs[i][5]:.2f}", (boxs[i][0],boxs[i][1]), cv2.FONT_HERSHEY_SIMPLEX, 1, (0,255,0), 2)
    return img


def resize_img(img):
     img = cv2.resize(img, (64, 128))
     return img

root_dir = r"F:\SY_data_set\晋云\20250102\source\fire"

src_dir = os.path.join(root_dir,"images")
xml_dir = os.path.join(root_dir,"voc_xml")
dst_dir = os.path.join(root_dir,"draw_img")
if not os.path.exists(dst_dir):
    os.makedirs(dst_dir)

xml_path = glob(xml_dir + r"/*.xml")

for i in range(len(xml_path)):
    print(xml_path[i])
    path_name = os.path.split(xml_path[i])[1]
    suffix = os.path.split(xml_path[i])[1].split(".")[0]
    for img_f in img_format:
        if os.path.exists(src_dir + "/" +path_name.replace(".xml",f".{img_f}")):
            img_path = src_dir + "/" +path_name.replace(".xml",f".{img_f}")
                    # if  img_path:
    # img_path = src_dir + "/" +path_name.replace(".xml",".jpg")
#     img = cv2.imread(img_path)
    img = cv2.imdecode(np.fromfile(img_path,dtype=np.uint8),1)
    tree = ET.parse(xml_path[i])
    root = tree.getroot()
    size = root.find('size')
    w = int(size.find('width').text)
    h = int(size.find('height').text)
    j = 0
    bboxes = []
    for obj in root.iter('object'):
        j += 1
        difficult = obj.find('difficult').text
        xmlbox = obj.find('bndbox')
        cls = obj.find('name').text

        # print(xmlbox.find('xmin').text.split(".")[0])
        b = [int(xmlbox.find('xmin').text.split(".")[0]),int(xmlbox.find('ymin').text.split(".")[0]) , 
             int(xmlbox.find('xmax').text.split(".")[0]),int(xmlbox.find('ymax').text.split(".")[0])]
        print("b:",b)
        b.append(cls)
        conf = random.uniform(0.6, 0.85)
        b.append(conf)
        bboxes.append(b)
        print(bboxes)

#         Rectimg = img[b[2]:b[3],b[0]:b[1]]
# #         cv2.imwrite(dst_dir + "/" + suffix + "_" + str(j) + ".jpg" ,Rectimg)
#         print(Rectimg.shape[0])
#         if Rectimg.shape[0] == 0:
#             continue
        # Rectimg = resize_img(Rectimg)
    draw_img = drow_rect(img,bboxes)
    # print(draw_img.shape)
    # cv2.imwrite("aaaa.jpg",draw_img)
    # save_path = dst_dir + "/" + cls 
    if not os.path.exists(dst_dir):
        os.makedirs(dst_dir)
    cv2.imencode('.jpg', draw_img)[1].tofile(dst_dir + "/" + suffix + ".jpg" )

            
print("OKK")

